package com.alatus.djl.service.impl;

import ai.djl.MalformedModelException;
import ai.djl.Model;
import ai.djl.basicdataset.cv.classification.ImageFolder;
import ai.djl.inference.Predictor;
import ai.djl.metric.Metrics;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.transform.Resize;
import ai.djl.modality.cv.transform.ToTensor;
import ai.djl.modality.cv.translator.ImageClassificationTranslator;
import ai.djl.ndarray.types.Shape;
import ai.djl.training.DefaultTrainingConfig;
import ai.djl.training.EasyTrain;
import ai.djl.training.Trainer;
import ai.djl.training.dataset.RandomAccessDataset;
import ai.djl.training.evaluator.Accuracy;
import ai.djl.training.listener.TrainingListener;
import ai.djl.training.loss.Loss;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import com.alatus.djl.Models;
import com.alatus.djl.service.TrainService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.BufferedWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

@Service
@Slf4j
public class TrainServiceImpl implements TrainService {
    @Override
    public void train(String datasetPath, String modelName, String modelPath) throws IOException, TranslateException {
//        获取数据集
        ImageFolder imageFolder = initDataset(Paths.get(datasetPath));
//        分隔数据集
//        这可以把图片数据集做切分,80%的数据用于训练,20%的数据用于测试
        RandomAccessDataset[] randomAccessDatasets = imageFolder.randomSplit(8, 2);
//        训练集
        RandomAccessDataset trainDatasets = randomAccessDatasets[0];
//        验证集
        RandomAccessDataset validationDatasets = randomAccessDatasets[1];
//        定义模型(对多层神经网络进行封装,封装数学函数)
        try(Model model = Models.getModel();){
//            模型训练器,模型训练配置
//            首先准备好我们的训练配置,然后获取训练器
//            训练配置需要我们传递损失函数
            DefaultTrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
                    .addEvaluator(new Accuracy())//计算精度
                    .addTrainingListeners(TrainingListener.Defaults.logging());//训练监听器,我们这里想要训练日志输出
//            获取训练器
            Trainer trainer = model.newTrainer(config);
            trainer.setMetrics(new Metrics());
//            初始化使用一张图片,三种RGB色彩,100*100像素(先启动)
            trainer.initialize(new Shape(1, 3, 100, 100));
//            训练,使用快速训练
//            使用刚刚生成的训练器进行拟合训练
            EasyTrain.fit(trainer, 2, trainDatasets, validationDatasets);
//            训练完成保存模型
            model.save(Paths.get(modelPath), modelName);
            List<String> synset = imageFolder.getSynset();
//            获取模型保存路径
            Path modelDir = Paths.get(modelPath);
//            直接在同一个目录下创建一个synset.txt文件
            Path resolve = modelDir.resolve("synset.txt");
            try(BufferedWriter bufferedWriter = Files.newBufferedWriter(resolve)){
//                把可迭代集合内容转为一个以换行符分隔的字符串,然后写入文件
                bufferedWriter.write(String.join("\n", synset));
            }
        };
        log.info("训练完成");
    }

    @Override
    public String predict(MultipartFile image) throws MalformedModelException, IOException, TranslateException {
        Image predictImage = ImageFactory.getInstance().fromInputStream(image.getInputStream());
//        拿到模型,算法模型只是一个公式,但是模型公式中每一个变量最终应该是多少都是训练得到的
        try(Model model = Models.getModel()){
//            这里的模型加载的本质就是我们训练得到的那一大堆的参数,所以加载模型就是把参数加载到内存中
//            再结合模型和公式,模型加载完成之后,模型就可以开始预测了
            model.load(Paths.get("build/models"));
//            模型是什么?
//            封装的计算函数和算法公式加上训练得到的一大堆参数数据,合在一起就是我们所称的模型
//            未训练的模型就是封装起来的算法公式和计算规则
//            训练好的模型就是公式和训练得到的一大堆参数数据,但是它本身没有推测作用,需要结合算法公式,才能进行预测

//            预测器里面需要一个转换器,来把需要预测的数据转换为模型认识可以理解的数据
            ImageClassificationTranslator classificationTranslator = ImageClassificationTranslator.builder()
                    .addTransform(new Resize(100, 100))//调整尺寸
                    .optApplySoftmax(true)
                    .addTransform(new ToTensor())//将图像N维数组从预处理的格式转换为神经网络格式的变换
                    .build();
//            模型需要一个预测器
            try(Predictor<Image, Classifications> predictor = model.newPredictor(classificationTranslator);){
                Classifications predict = predictor.predict(predictImage);
                return predict.toString();
            }
        }
    }

    private ImageFolder initDataset(Path path) throws IOException, TranslateException {
//        当数据集完成加载以后,数据集就知道我们的图片要分成几类了,因为它是根据文件夹加载的
        ImageFolder folder = ImageFolder.builder().setRepositoryPath(path)//指定数据集的位置
                .setSampling(128, true)
                .optMaxDepth(10)
                .addTransform(new Resize(100, 100))
                .addTransform(new ToTensor())//把图片转为向量
                .build();
        folder.prepare(new ProgressBar());
        List<String> synset = folder.getSynset();
        for (String s : synset) {
            log.info(s);
        }
//        之前的数字分类案例我们保存了一个synset的文件,这个文件就是用来保存图片的类别的枚举值
//        顺手保存一下syset文件也可以
        return folder;
    }
}
